import pandas as pd
import urllib
import time
import sys

#script that loops over all the public companies in allpubliccomp.csv, extracts the financials using YQL then store the
#the results in a new csv file

#note the financials are annual

count = 1
total = 1

def financial_variables(stock):
    #function to query Yahoo YQL and return the financial data for a stock based on stock symbol.
    #the variables are selected according to Altman Z-Score
    stock = stock.strip('$')
    total_assets=""
    total_liability=""
    current_assets=""
    current_liability=""
    retained_earnings=""
    market_capital = ""
    ebitda=""
    sales=""
    stockprice=""


    try:
        baseurl = "https://query.yahooapis.com/v1/public/yql?"
        #extract data from balance sheet
        yql_bs_query = "select * from yahoo.finance.balancesheet where symbol in ('"+stock+"') and timeframe='annual'"
        yql_bs_url = baseurl + urllib.parse.urlencode({'q':yql_bs_query}) + "&format=json&diagnostics=true&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys&callback="
        bs_json = pd.io.json.read_json(yql_bs_url)
        if bs_json["query"]["results"]["balancesheet"]["statement"] is not None:
            try:
                total_assets=bs_json["query"]["results"]["balancesheet"]["statement"][0]['TotalAssets']['content']
            except:
                total_assets="0"
            try:
                total_liability=bs_json["query"]["results"]["balancesheet"]["statement"][0]['TotalLiabilities']['content']
            except:
                total_liability="0"
            try:
                current_assets=bs_json["query"]["results"]["balancesheet"]["statement"][0]['TotalCurrentAssets']['content']
            except:
                current_assets="0"
            try:
                current_liability=bs_json["query"]["results"]["balancesheet"]["statement"][0]['TotalCurrentLiabilities']['content']
            except:
                current_liability = "0"
            try:
                retained_earnings=bs_json["query"]["results"]["balancesheet"]["statement"][0]['RetainedEarnings']['content']
            except:
                retained_earnings="0"
             #extract data from income statement
            yql_is_query = "select * from yahoo.finance.incomestatement where symbol in ('"+stock+"') and timeframe='annual'"
            yql_is_url = baseurl + urllib.parse.urlencode({'q':yql_is_query}) + "&format=json&diagnostics=true&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys&callback="
            is_json = pd.io.json.read_json(yql_is_url)

            try:
                ebitda=is_json["query"]["results"]["incomestatement"]["statement"][0]['EarningsBeforeInterestAndTaxes']['content']
            except:
                ebitda="0"
            try:
                sales=is_json["query"]["results"]["incomestatement"]["statement"][0]['TotalRevenue']['content']
            except:
                sales="0"

            #extract data from finance quotes
            yql_qt_query = "select * from yahoo.finance.quotes where symbol in ('"+stock+"')"
            yql_qt_url = baseurl + urllib.parse.urlencode({'q':yql_qt_query}) + "&format=json&diagnostics=true&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys&callback="
            qt_json = pd.io.json.read_json(yql_qt_url)
            try:
                stockprice = qt_json["query"]["results"]["quote"]["LastTradePriceOnly"]
            except:
                stockprice = "0"
            try:
                market_capital = qt_json["query"]["results"]["quote"]['MarketCapitalization']
            except:
                market_capital = "0"

    except:
        print ("error on ", stock, ". Error message",sys.exc_info())
    global count
    count = count +1
    global total
    print ("\n",count,"/",total,":",stock, "|", total_assets, "|", total_liability, "|", current_assets,"|",current_liability,"|",retained_earnings,"|",market_capital,"|",ebitda,"|",sales,"|", stockprice)
    return ([total_assets, total_liability, current_assets,current_liability,retained_earnings,market_capital,ebitda, sales, stockprice])

def run():
    firms_fin_location = "data/allpubliccomp.csv"
    firms = pd.read_csv(firms_fin_location)
    total = len(firms.index)
    firms['altman_variables']=firms['Symbol'].apply(financial_variables)
    firms.to_csv('data/altman_results4.csv')

financial_variables('tfsc')